from datetime import datetime, timedelta
from random import choice, randint, sample
import numpy as np
import pandas as pd
from pymongo import MongoClient
from collections import defaultdict

# 1. 行为权重体系（总和不需要为1，后面会归一化）
BEHAVIOR_WEIGHTS = {
    'click': 0.1,  # 点击
    'read': 2.0,  # 阅读（提高比重）
    'review': 2.5,  # 评论（提高比重）
    'note': 1.5,  # 笔记
    'like': 0.4,  # 点赞
    'share': 1.0,  # 分享
    'collect': 1.8  # 收藏（提高比重）
}

# 评论分数范围
REVIEW_SCORE_RANGE = (1, 5)  # 1-5分


# 2. 从 MongoDB 获取所有图书 ID
def load_book_ids_from_mongo():
    try:
        client = MongoClient('mongodb://localhost:27017/')
        db = client['test']
        collection = db['book']
        all_book_ids = [str(doc['_id']) for doc in collection.find({}, {'_id': 1})]

        print(f"已获取{len(all_book_ids)}本图书作为候选")
        return all_book_ids
    except Exception as e:
        print(f"加载失败，使用默认 ID: {e}")
        return []


# 3. 生成日志数据（核心功能）
def generate_logs(book_ids, num_records=20000):
    """生成大规模行为日志，包含评论分数"""
    records = []
    behavior_list = list(BEHAVIOR_WEIGHTS.keys())

    # 计算权重总和并归一化
    total_weight = sum(BEHAVIOR_WEIGHTS.values())
    normalized_weights = [w / total_weight for w in BEHAVIOR_WEIGHTS.values()]

    # 按权重随机选择行为
    actions = np.random.choice(
        behavior_list,
        size=num_records,
        p=normalized_weights
    )

    # 预先生成评论分数（仅用于review行为）
    review_scores = np.random.randint(*REVIEW_SCORE_RANGE, size=num_records)

    for i in range(num_records):
        book_id = choice(book_ids)
        action = actions[i]

        record = {
            'u_id': randint(1, 100),
            'b_id': book_id,
            'action': action,
            'timestamp': datetime.now() - timedelta(hours=np.random.uniform(0, 30 * 24))
        }

        if action == 'review':
            record['score'] = review_scores[i]  # 添加评论分数

        records.append(record)

    # 计算统计信息
    book_counts = defaultdict(int)
    action_counts = defaultdict(int)
    for r in records:
        book_counts[r['b_id']] += 1
        action_counts[r['action']] += 1

    print("\n日志生成统计：")
    print(f"- 总记录数: {len(records)}")
    print(f"- 涉及图书数: {len(book_counts)}")
    print(f"- 行为分布:")
    for action, count in action_counts.items():
        print(f"  {action}: {count}次 ({count / len(records) * 100:.1f}%)")

    return pd.DataFrame(records)


# 4. 主函数（仅生成日志）
if __name__ == "__main__":
    # 连接MongoDB
    client = MongoClient('mongodb://localhost:27017/')
    db = client['test']
    logs_collection = db['logs']
    books_collection = db['book']

    # 加载图书ID
    book_ids = load_book_ids_from_mongo()

    # 生成日志数据
    log_df = generate_logs(book_ids, num_records=20000)

    # 保存到MongoDB
    try:
        logs_collection.delete_many({})
        logs_collection.insert_many(log_df.to_dict(orient='records'))
        print(f"\n成功插入 {len(log_df)} 条日志记录到MongoDB")
    except Exception as e:
        print(f"插入MongoDB失败: {e}")

    # 保存到CSV
    log_df.to_csv('book_logs.csv', index=False, encoding='utf_8_sig')
    print("日志已保存到 book_logs.csv")